#include <stdio.h>
#include <stdarg.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <ctype.h>
#include <time.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <errno.h>
#include <unistd.h>
#include <assert.h>
#include <netinet/in.h>
#include <netinet/in_systm.h>
#include <netinet/ip.h>
#include <netdb.h>
#include <signal.h>

#include "datalink.h"
#include "kip.h"
#include "ktcp.h"
#include "log.h"

#include "kmgmt.h"

#if defined (HAVE_DMALLOC_H) && defined (HAVE_LIBDMALLOC)
#include "dmalloc.h"
#endif
enum {FALSE,TRUE};
/* Maximum Payload size is a full size KTCP packet */
#define MAX_PAYLOAD_SIZE	20+536
// #define MAX_PAYLOAD_SIZE	20+236
#define	KIP_HEADER_SIZE		20
#define MAX_PACKET_SIZE		(KIP_HEADER_SIZE+MAX_PAYLOAD_SIZE)

#define MAX_NUM_SOCKFDS		4

/* Maximum Transmission Unit */
#define MTU	576

#define KIP_DEFAULT_TIMEOUT_TTL 10
#define KIP_TIME_TIMEOUT_PERIOD 2000

static int slow_timeout = 1;
static unsigned short ip_id;
route * last_route;

static list ip_reass_ctx_list;

typedef struct ip_context_t {
	ushort id;		//ip_id
	struct in_addr src;  //src;
	list fragments; 	//list received fragments of the ctx
	bool end,done; 	//last fragment. not received -> null
	ushort total, recv;
} * ip_context;

typedef struct ip_fragment_t{
	struct ip * header;	//header
	ushort offset;		//offset
	ushort length;		//length of data
	void * data;		//data
} * ip_fragment;


void ip_dump_header(char *title,struct ip *);
/**
 * initialize KIP layer. loading routing table and firewall configuration
 * should be done in this function. This function is called before KENS
 * enters its main dispatch loop.
 * @param conf a pointer to KENS configuration file hash structure
 * @return 0 when success, others when errors
 */
int ip_init(dictionary *conf)
{
	ip_id ^= (unsigned char)((getpid()*getppid()*time(NULL)) & 0xffff);

	rt_init(conf);

	last_route = (route*) malloc(sizeof(route));
	last_route=malloc(sizeof(route));
	inet_aton("0.0.0.0",&last_route->ro_dst);
	ip_reass_ctx_list = list_open();
	L_IP("Start logging ip");

	return 0;
}

/**
 * shutdown and free whole resources which were allocated at ip_init() or
 * during KENS runtime, especially resources related to routing table and 
 * firewall. This function is called after KENS receive a termination signal
 * from user.
 * @return 0 when success, others when errors
 */
int ip_shutdown(void)
{
	rt_cleanup();
	return 0;
}
void ip_remove_finished_context()
{
	ip_context ctx = NULL;
	ip_fragment frag = NULL;
	while ((ctx = list_get_head(ip_reass_ctx_list)) && ctx->done) {
		// L_IP("Time tick ip_reass_ctx_list has %d contexts",list_size);
		L_IP("Delete a context from %s id=%d",inet_ntoa(ctx->src),ctx->id);
		while (frag=list_remove_head(ctx->fragments)) {
			free(frag->header);
			free(frag->data);
			free(frag);
		}
		list_close(ctx->fragments);
		list_remove(ip_reass_ctx_list,ctx);
		free(ctx);
	}
	
}
/** 
 * This fuction reduces the ttl of the fragment packets in reassembly list. 
 * Also, it checks the timeout.
 */
void ip_slow_timeout() {

	if ((++slow_timeout % KIP_TIME_TIMEOUT_PERIOD) != 0)
		return;
	int list_size = list_get_count(ip_reass_ctx_list);
	ip_context ctx = NULL;
	ip_fragment frag = NULL;
	list_position ctx_pos = NULL;
	ctx_pos=list_get_head_position(ip_reass_ctx_list);
	while (ctx_pos!=NULL) {
		ctx = list_get_at(ctx_pos);
		if (!ctx->done){
			frag=list_get_head(ctx->fragments);
			--frag->header->ip_ttl;
			if (frag->header->ip_ttl == 0) {
				L_IP("(%08x) Context from %s id=%d timeout",ctx,inet_ntoa(ctx->src),ctx->id);
				ctx->done = true;			
			}
		}
		ctx_pos = list_get_next_position(ctx_pos);
	}
	ip_remove_finished_context();
}

/**
 * this function is called by KENS_dispatch_loop() in Kernel_main.c .
 * Nothing has to be more inplemented in this function.
 */
int ip_dispatch() {
	ip_slow_timeout();
	return 0;
}

ushort ip_cksum(struct ip *ip,ushort len)
{
	long sum =0;
	ushort *temp = (ushort *)ip;
	ushort old_sum = ip->ip_sum;
	unsigned char old_ttl = ip->ip_ttl;
	ip->ip_sum = 0;
	ip->ip_ttl = 0;
	while (len>1) {
		sum += *(temp++); 
		if (sum & 0x80000000)
			sum = (sum & 0xFFFF) + (sum >>16);
		len -= 2;
		// L_IP("	=%u",sum);
	}
	if (len)
		sum += (ushort) *(unsigned char *)temp;
	while (sum >> 16) 
		sum = (sum & 0xFFFF) + (sum >>16);
	ip->ip_sum = old_sum;
	ip->ip_ttl = old_ttl;
	return ~sum;
}
 

ip_context find_ctx(ushort id,struct in_addr src, list l) {

	//find the context having that id, inside the list of context
	ip_context ctx = NULL; 
	list_position pos = NULL;
	for (pos=list_get_head_position(l);
			pos!=NULL; pos = list_get_next_position(pos)) {
		ctx = list_get_at(pos);
		if (ctx->id==id && ctx->src.s_addr == src.s_addr)
			break;
	}

	if (pos==NULL) {
		L_IP("Ctx with id %d not found, create new ctx", id);
		ctx=malloc(sizeof(struct ip_context_t ));
		ctx->id=id;
		ctx->src.s_addr = src.s_addr;
		ctx->end=FALSE;
		ctx->done = FALSE;
		ctx->total = ctx->recv = 0;
		ctx->fragments = list_open();
		list_add_head(l,ctx);
		L_IP("Reass list have %d contexts", list_get_count(l));
	}
	assert(ctx->id==id && ctx->src.s_addr == src.s_addr);
	return ctx; 
}

#define pos2off(X) (((ip_fragment) list_get_at(X) )->offset)

void insert_frag(ip_context ctx,ip_fragment fragment) {
	list_position iter=NULL;
	ushort offset = fragment->offset;
	//create ctx entry if necessary
	if (list_get_count(ctx->fragments)== 0) {
		list_add_head(ctx->fragments,fragment);
	} else {
		// try to put fragment in ctx->fragments, keep the order
		iter = list_get_tail_position(ctx->fragments);
		while (iter && offset < pos2off(iter))
			iter = list_get_prev_position(iter);
		if (!iter)
			list_add_head(ctx->fragments,fragment);
		else {
			list_insert_after(iter,fragment);
			// L_IP("insert before %d",pos2off(iter));
		}
	}
	// L_IP("count = %d",list_get_count(ctx->fragments));
}


/**
 * This fuction is called by ip_input function when the incoming packet is
 * a fragment packet. 
 * You should add the fragment to a proper locaiton in reassembly list.
 * If reassembly is possible, reassemble the fragments.
 * @return 0 when success.
 */
int ip_reass(void *buf) {
	struct ip * hdr = (struct ip *) buf;
	ushort id = hdr->ip_id;
	ushort hlen = sizeof(struct ip);
	ushort dlen = hdr->ip_len-hlen;
	ushort offset = hdr->ip_off & IP_OFFMASK;

	ip_fragment fragment = malloc(sizeof(struct ip_fragment_t));
	fragment->length = dlen;
	fragment->offset = offset;

	fragment->header = malloc(hlen);
	memcpy(fragment->header,buf,hlen);

	fragment->data = malloc(dlen);
	memcpy(fragment->data,buf + hlen,dlen);

	ip_context ctx = find_ctx(id,hdr->ip_src,ip_reass_ctx_list);
	L_IP("(%08x) get fragment at %d/%d",ctx,offset,ctx->total);
	insert_frag(ctx,fragment);

	//MF flags is turn off (fragment with last offset)
	if (!(hdr->ip_off & IP_MF))
	{
		ctx->end=TRUE;
		ctx->total = fragment->length + fragment->offset;
	}

	ctx->recv += fragment->length;
	L_IP("(%08x) received %d/%d",ctx,ctx->recv,ctx->total);
	assert(!ctx->end || ctx->recv <= ctx->total);
	//received all the parts
	if (ctx->end && ctx->recv==ctx->total) {
		L_IP("(%08x) Got all fragment",ctx);
		char *buf1 = malloc(ctx->total);
		struct ip ip = *fragment->header;
		list_position fragment_pos = NULL;
		offset = 0;
		fragment_pos = list_get_head_position(ctx->fragments);
		while (fragment_pos){
			fragment = list_get_at(fragment_pos);
			L_IP("(%08x) reassemble fragment %d=%d/%d",ctx,offset,fragment->offset,ctx->total);
			assert(offset == fragment->offset);
			memcpy(buf1+offset,fragment->data,fragment->length);
			offset += fragment->length;
			fragment_pos = list_get_next_position(fragment_pos);
		}
		ctx->done = TRUE;
		assert(offset == ctx->total);
		tcp_dispatch_in(ip.ip_src,ip.ip_dst,buf1,ctx->total);
		free(buf1);
		ip_remove_finished_context();
		return 0;
	}
	return 0;
}


int ip_fwd(void *buf,int len) {
	L_IP("Start ip_forward");
	struct ip *ip;
	ip = (struct ip *)buf;

	int error, type=0, code;
	struct rtentry_t *rt;

	if (ip->ip_ttl <= 1) { L_IP("IP Address List are NULL"); return 0; }
	ip->ip_ttl -= 1;
	L_IP("Decrease TTL");


	error = ip_output2(ip->ip_src,ip->ip_dst,buf,len,last_route);
}

/**
 * this function is called for every incoming packet from datalink.
 * You should implement logic in this function or in serveral functions
 * to process incoming packets to this host as well as packets which
 * pass through this host, routing or forwarding.
 *
 * For further information, I strongly recommend you to refer TCP/IP
 * Illustrated Volume 2.
 *
 * @param buf buffer which contains data from datalink. this buffer starts
 * from KIP header
 * @param len exact size of data contained in buf
 * @return 0 when success, others when errors
 */

#define R_ERR -1

int ip_input(void *buf,int len)
{
	L_IP("Start ip_input");
	struct ip *ip;
	struct in_ifaddr_t * ia;
	unsigned short hlen,ip_len;
	if (!buf) { L_IP("Buf is null"); return R_ERR; }

	ip = (struct ip *)buf;

	if (_ifnet==NULL) { L_IP("IP Address List are NULL"); return R_ERR; }

	ip_len = ntohs((unsigned short)ip->ip_len) & IP_OFFMASK;
	if (ip_len<KIP_HEADER_SIZE) { L_IP("ip_len Invalid"); return R_ERR; }

	hlen = ip->ip_hl << 2;
	if (hlen<KIP_HEADER_SIZE) { L_IP("hlen Invalid"); return R_ERR; }
	if (hlen>ip_len) { L_IP("Header length Invalid"); return R_ERR; }

	if (ip_len<len) { L_IP("IP Packet length smaller than buffer length"); len=ip_len; }

	if (ip_len>len)  { L_IP("IP Packet length bigger than buffer length"); return R_ERR; }

	// Check version
	if (ip->ip_v!=4) { L_IP("Version in correct!"); return R_ERR; }

	// Host convertion
	ip->ip_len=ip_len;
	ip->ip_id=ntohs((unsigned short)ip->ip_id);
	ip->ip_off=ntohs((unsigned short)ip->ip_off);

	ip_dump_header("",ip);
	// Checksum
	if (ip->ip_sum != ip_cksum(ip,len)) {
		L_IP("Checksum fail!");
		return R_ERR;
	}

	// check address

	for (ia = _in_ifaddr; ia; ia=ia->ia_next) {
		if (ia->ifa->ifa_addr.s_addr == ip->ip_dst.s_addr) {
			/* now it is ready to send to KTCP layer */
			L_IP("Destination is our address");

			if (ip->ip_off & ~IP_DF) {
				//offset turn on or MF turn on-> reass
				L_IP("Received a fragment");
				return ip_reass(buf);
			}
			L_IP("Received a full datagram");
			buf+=sizeof(struct ip);
			tcp_dispatch_in(ip->ip_src,ip->ip_dst,buf,ip_len-KIP_HEADER_SIZE);
			return 0;
		}
	}

	//forwarding;
	ip_fwd(buf,len);

	return 0;
}

int ip_output2(struct in_addr src,struct in_addr dst,
		void *buf,size_t len,route *ro)
{
	L_IP("Start ip_output2");
	struct ip *ip = (struct ip *)buf;
	ifnet *ifp = NULL;

//	ifp = ifunit("lo");

	//select route
	if (!ro) ro=last_route;
	if (!ro || !ro->ro_rt || last_route->ro_dst.s_addr!=ip->ip_dst.s_addr) {

		// L_IP("Allocate route");
		last_route->ro_dst=ip->ip_dst;

		L_IP("Call rt_alloc");
		rt_alloc(last_route);
		ro=last_route;
		if (!ro) {
			L_IP("No route found!"); 
			return -1; 
		} else {
			L_IP("Foudn a route!"); 
		}
	} else {
		L_IP("Use old route");
	}

#define ifatoia(ifa) ((in_ifaddr *) ifa)
#define sintosa(sin) ((struct sockaddr *) (sin))

//	in_ifaddr * ia = ifatoia(ro->ro_rt->rt_ifa);
	ifp = ro->ro_rt->rt_ifp;
	ro->ro_rt->rt_use++;

//	L_IP("Can reach here!");
	if (ro->ro_rt->rt_flags & RTF_GATEWAY)
		dst = ro->ro_rt->gw;
	ip->ip_len = htons((unsigned short)ip->ip_len);
	ip->ip_off = htons((unsigned short)ip->ip_off);
	ip->ip_id = htons(ip->ip_id);


//	if (ip->ip_src.s_addr == INADDR_ANY)
//		ip->ip_src = IA_
	// L_IP("Now send to dl");
	dl_output(ifp,buf,len,dst);

	return 0;
}

/**
 * this function deals outgoing packet from this host. That is,
 * this function is directly called by Transport layer, KTCP.
 * You should implement logic to decide packet route which is sometimes
 * given by KTCP layer, fragment data and generate proper KIP header
 * for each outgoing packet. You may leave this function as a wrapper
 * to another extended KIP output function which receive an extra argument
 * ,flag, to share the facilities with packet forwarding.
 *
 * For further information, I strongly recommend you to refer TCP/IP
 * Illustrated Volume 2.
 *
 * @param src where this packet comes from in network byte order.
 * @param dst where this packet goes to in network byte order.
 * @param buf packet's payload
 * @param len length of packet's payload
 * @param ro routing information which is already cached by previous call.
 * If this argument is not null and the structure is empty, ip_output will
 * fill the structure with proper routing information for next time use.
 * If this argument is not null and the structure is not empty, ip_output will
 * use given routing information without lookup routing table.
 * @return 0 when success, others when errors
 */
int ip_output(struct in_addr src,struct in_addr dst,
		void *buf,size_t len,route *ro)
{
	L_IP("Start ip_output");
	unsigned char *pkt = buf;
	int hlen = sizeof(struct ip);
	int dlen;
	struct ip *ip;
	ifnet *ifp = NULL;
	unsigned char *ip_buf;
	int offset;

	ip_buf = malloc(MAX_PACKET_SIZE+len);
	memcpy(ip_buf+KIP_HEADER_SIZE,pkt,len);	/* copy TCP packet into local buffer */
	ip = (struct ip *)ip_buf;	/* may grab first 20 bytes */

	ip->ip_v = 4;	/* version 4 */
	ip->ip_hl = hlen >> 2;
	ip->ip_ttl = KIP_DEFAULT_TTL;

	ip->ip_src = src;
	ip->ip_dst = dst;


	for (offset = 0;offset < len;offset += MAX_PAYLOAD_SIZE)
	{
		ip->ip_off = offset;
		if (len <= MAX_PAYLOAD_SIZE)
			L_IP("Full segment");
		else
			L_IP("a fragment");
		ip->ip_len = len - offset;
		if (ip->ip_len > MAX_PAYLOAD_SIZE) {
			ip->ip_len = MAX_PAYLOAD_SIZE;
			ip->ip_off |= IP_MF;
		} 
		dlen = (ip->ip_len & IP_OFFMASK);
		ip->ip_len += KIP_HEADER_SIZE;
		ip->ip_id = ip_id;			/* Fragments have the same id */
		memcpy(ip_buf + KIP_HEADER_SIZE,buf+offset,dlen);
		ip->ip_sum = ip_cksum(ip,ip->ip_len);

		ip_dump_header("",ip);
		L_IP("Send to ip_output2 %d/%d ",offset + dlen,len);

		// dl_output(ifp,ip_buf,hlen+len,dst);
		//'cause it need to find the route also - ipoutput2 find route
		ip_output2(src, dst, ip_buf, dlen+KIP_HEADER_SIZE, ro);
	}

	L_IP("end ip_output");

	ip_id++;
	free(ip_buf);
	return 0;
}

/**
 * get virtual ip address of virtual network interface card
 * which is used to transmit a packet destined to 'in'.
 * this is called by TCP layer to decide source ip address
 * for outgoing packet. Thus, you should consult routing table
 * to figure out which device should be used.
 * @param in destination of a packet in network byte order
 * @return source ip address of interface card in network byte order
 */
uint32_t ip_host_address(struct in_addr in)
{
	route ro;	
	ifnet *ifp;

	for ( ifp = _ifnet; ifp; ifp = ifp->if_next )
		if ( in_localnet(ifp,in) )
			return IN_ADDR(ifp).s_addr;

	ro.ro_dst.s_addr = in.s_addr;
	ro.ro_rt = NULL;

	rt_alloc(&ro);

	if ( ro.ro_rt == NULL )
		return htonl(INADDR_ANY);

	return IN_ADDR(ro.ro_rt->rt_ifp).s_addr;
}

/**
 * dump given KIP header into log file.
 * @param title a header for each log
 * @param ip a KIP header to dump
 */
void ip_dump_header(char *title,struct ip *ip)
{
	char buf[80];

	inet_ntoa_r(ip->ip_src,buf,80);

	L_IP_HDR("%s v:%d src:%s dest:%s ttl:%d",
			title,
			ip->ip_v,
			buf,
			inet_ntoa(ip->ip_dst),
			ip->ip_ttl
	);
	L_IP_HDR("id=%d off=%d len=%d",ip->ip_id,ip->ip_off & IP_OFFMASK, ip->ip_len);
}

